import pandas as pd
import pymysql
from sqlalchemy import create_engine
from sqlalchemy.exc import SQLAlchemyError
import argparse


def read_excel_file(file_path):
    """读取Excel文件并返回DataFrame"""
    try:
        df = pd.read_excel(file_path)
        print(f"成功读取Excel文件，共{len(df)}行，{len(df.columns)}列")
        return df
    except Exception as e:
        print(f"读取Excel文件失败: {e}")
        raise


def prepare_data_for_import(df, table_columns):
    """准备数据以导入数据库，确保列名和数据类型匹配"""
    # 获取DataFrame的列名，并转为小写（假设数据库表列名是小写的）
    df_columns = [col.lower() for col in df.columns]

    # 检查DataFrame的列是否与数据库表的列匹配
    missing_columns = set(table_columns) - set(df_columns)
    extra_columns = set(df_columns) - set(table_columns)

    if missing_columns:
        print(f"警告: Excel文件中缺少这些列: {missing_columns}")

    if extra_columns:
        print(f"警告: Excel文件中包含这些额外的列，将被忽略: {extra_columns}")

    # 选择并重新排序DataFrame的列以匹配数据库表
    valid_columns = [col for col in table_columns if col in df_columns]
    df = df[[col for col in df.columns if col.lower() in valid_columns]]

    # 重命名DataFrame的列以匹配数据库表
    df.columns = [col.lower() for col in df.columns]

    # 处理日期列
    for column in df.columns:
        if pd.api.types.is_datetime64_any_dtype(df[column]):
            df[column] = df[column].dt.strftime('%Y-%m-%d %H:%M:%S')

    # 根据date列填充year、month、day和season列
    if 'date' in df.columns:
        df['date'] = pd.to_datetime(df['date'])
        df['year'] = df['date'].dt.year
        df['month'] = df['date'].dt.month
        df['day'] = df['date'].dt.day
        df['season'] = pd.cut(df['date'].dt.month, bins=[0, 3, 6, 9, 12], labels=['spring', 'summer', 'fall', 'winter'])

    return df


def get_table_columns(db_config, table_name):
    """获取数据库表的列名"""
    conn = pymysql.connect(
        host=db_config['host'],
        user=db_config['user'],
        password=db_config['password'],
        database=db_config['database'],
        charset='utf8mb4'
    )

    try:
        with conn.cursor() as cursor:
            cursor.execute(f"SHOW COLUMNS FROM {table_name}")
            columns = [row[0] for row in cursor.fetchall() if row[0]!= 'id']  # 排除自动递增的id列
            return columns
    finally:
        conn.close()


def import_data_to_mysql(df, table_name, db_config):
    """将数据导入到MySQL数据库"""
    try:
        # 创建SQLAlchemy引擎
        engine = create_engine(
            f"mysql+pymysql://{db_config['user']}:{db_config['password']}@{db_config['host']}/{db_config['database']}?charset=utf8mb4"
        )

        # 将数据导入数据库，使用append模式
        df.to_sql(table_name, con=engine, if_exists='append', index=False)
        print(f"成功导入 {len(df)} 行数据到表: {table_name}")

    except SQLAlchemyError as e:
        print(f"数据库操作错误: {e}")
        raise
    except Exception as e:
        print(f"导入数据失败: {e}")
        raise


def main():
    parser = argparse.ArgumentParser(description='将Excel数据导入MySQL数据库')
    parser.add_argument('--file', required=True, help='Excel文件路径')
    parser.add_argument('--table', required=True, help='数据库表名')
    parser.add_argument('--host', default='localhost', help='数据库主机地址')
    parser.add_argument('--user', default='root', help='数据库用户名')
    parser.add_argument('--password', required=True, help='数据库密码')
    parser.add_argument('--database', required=True, help='数据库名')

    args = parser.parse_args()

    db_config = {
        'host': args.host,
        'user': args.user,
        'password': args.password,
        'database': args.database
    }

    try:
        # 读取Excel文件
        df = read_excel_file(args.file)

        # 获取数据库表的列名
        table_columns = get_table_columns(db_config, args.table)
        print(f"数据库表 {args.table} 的列: {table_columns}")

        # 准备数据
        prepared_df = prepare_data_for_import(df, table_columns)

        # 导入数据
        import_data_to_mysql(prepared_df, args.table, db_config)

    except Exception as e:
        print(f"程序执行失败: {e}")
        exit(1)


if __name__ == "__main__":
    main()

    # use
    # python only_import_xlsx_data.py --file ../data/new_data_2025_processed.xlsx --table weather_data --host localhost --user root --password 1026 --database guilin_temp